"""
Monte-Carlo Tree Search
"""
import math as m
import numpy as np

from lib import game, model

import torch.nn.functional as F


class MCTS:
    """
    Class keeps statistics for every state encountered during the search
    用来保存在对弈过程中预测到的接下来可能的未来操作的分支会遇到的奖励回报统计
    """
    def __init__(self, c_puct=1.0):
        # todo 作用 但是在本代码中并没有使用到
        # 用于节点的选择过程
        self.c_puct = c_puct
        # count of visits, state_int -> [N(s, a)]
        # 这个应该是当前状态s进行当前动作a的预测访问次数，
        # 该动作被预测的次数越多，则越会被采纳执行
        self.visit_count = {}
        # total value of the state's action, state_int -> [W(s, a)]
        # 当前状态s执行每个动作的奖励回报 todo 作用
        self.value = {}
        # average value of actions, state_int -> [Q(s, a)]
        # 计算当前状态s执行动作a的平均回报，即value / visit_count todo 作用
        self.value_avg = {}
        # prior probability of actions, state_int -> [P(s,a)]
        # 当前状态下，能够执行的动作，每一个动作的概率
        # 当前状态s执行动作a的概率
        self.probs = {}

    def clear(self):
        '''
        清理游戏状态缓存，在更新最佳模型时使用
        todo 什么时候使用
        '''

        self.visit_count.clear()
        self.value.clear()
        self.value_avg.clear()
        self.probs.clear()

    def __len__(self):
        '''
        获取当前存储状态、动作对应回报的数量
        '''
        return len(self.value)

    def find_leaf(self, state_int, player):
        """
        Traverse the tree until the end of game or leaf node
        搜索树，直到游戏结束或者没有发现未被探索的叶子
        在搜索的过程中，会记录访问过的状态和执行的动作，用于后续更新节点的统计信息
        :param state_int: root node state 搜索的起始起点
        :param player: player to move
        :return: tuple of (value, leaf_state, player, states, actions)
        1. value: None if leaf node, otherwise equals to the game outcome for the player at leaf 为0或者-1，表示是平局还是对手输了，为None表示是叶子节点，为数字表示是最终状态
        2. leaf_state: state_int of the last state 探索过程中，到达的最终状态
        3. player: player at the leaf node # 最终状态的时候，是哪个玩家的回合
        4. states: list of states traversed # 探索过程中，访问过的状态
        5. list of actions taken # 探索过程中，执行的动作（每个状态下执行的动作）

        执行树的一次遍历
        """
        # todo 作用
        states = [] # 存储访问过的游戏状态，由于是棋盘类游戏，所以所谓的游戏状态或者游戏过程节点，就是棋盘整体的状态
        actions = [] # 选出遍历过程中，根据当前的状态选择要执行的动作
        cur_state = state_int # 开始搜索的初始状态
        cur_player = player
        value = None # 存储搜索过程中，如果执行动作后，赢得了比赛，对手的得分（因为一旦下棋后，就轮到了对手的状态，所以value表示对手的得分）

        # 不是叶子节点则一直遍历
        while not self.is_leaf(cur_state):
            states.append(cur_state)

            # 提取当前状态访问的次数
            counts = self.visit_count[cur_state]
            # todo 计算所有元素的平方根
            total_sqrt = m.sqrt(sum(counts))
            # 提取当前状态下，能够执行的动作，每一个动作的概率
            # 当前游戏下，就是选择哪一列能下棋的概率
            probs = self.probs[cur_state]
            # 当前状态下，能够执行的动作，每一个动作的奖励回报Q值
            # todo values_avg什么时候更新的
            values_avg = self.value_avg[cur_state]

            # choose action to take, in the root node add the Dirichlet noise to the probs
            # 如果当前是根节点，则需要在根节点上选择要采取的动作，在概率上添加狄利克雷噪声。
            # todo 狄利克雷噪声是什么？
            if cur_state == state_int:
                # 创建一个狄利克雷噪声，用于在根节点上添加噪声
                noises = np.random.dirichlet([0.03] * game.GAME_COLS)
                # 在每个执行动作上添加噪音，进行探索 todo 看代码，怎么添加噪音
                probs = [0.75 * prob + 0.25 * noise for prob, noise in zip(probs, noises)]
            #评估并选择下一个扩展的节点。
            # 在MCTS的选择阶段，目标是从当前节点的所有子节点中选出一个进行进一步的探索。为了做出这个选择，每个子节点都会被赋予一个评分，这个评分基于它的潜在价值和探索的优先级
            # 计算过程如下：
            # value：代表该节点的平均价值，通常是从该节点开始的模拟/探索的历史结果的平均值
            # prob：这是从父节点到这个子节点的转移概率，通常由策略网络（如在AlphaGo或类似系统中）给出。 todo 什么时候更新的？
            # count：这是该节点的访问次数，用于控制探索的优先级。访问次数越多，则越会被采纳执行
            # total_sqrt：这是所有子节点访问次数的总和的平方根，用于调整探索与利用的平衡。
            # c_puct：这是一个超参数，用于控制探索与利用的平衡。
            # 这边还是计算了所有动作的分数，并未考虑到是否可以下棋
            score = [value + self.c_puct * prob * total_sqrt / (1 + count)
                     for value, prob, count in zip(values_avg, probs, counts)]
            # 从游戏的所有列索引中（也就是理论上所有可以执行的动作） 排除掉）
            invalid_actions = set(range(game.GAME_COLS)) - set(game.possible_moves(cur_state))
            # 如果哪一个动作是不能执行的，那么就将该动作的分数设置为负无穷
            for invalid in invalid_actions:
                score[invalid] = -np.inf
            # 选择分数最大的列索引，作为下一步要执行的动作
            action = int(np.argmax(score))
            # 将执行的动作存储到actions
            actions.append(action)
            # todo ，给该move方法添加注解
            # 游戏执行动作，返回游戏状态和是否胜利
            cur_state, won = game.move(cur_state, action, cur_player)
            if won:
                # if somebody won the game, the value of the final state is -1 (as it is on opponent's turn)
                # 如果游戏得到胜利，则当前状态的价值为-1，因为当前状态是对手的回合
                # 为什么当前的状态是对手的状态
                # 因为下棋后，轮到对手下棋，所以当前状态是对手的状态
                value = -1.0
            # 更新当前玩家
            cur_player = 1-cur_player
            # check for the draw
            # 如果当前value是None，表示还没得到胜利，那么就判断是否是平局
            # 如果当前状态下，所有的动作都不能执行，那么就是平局
            if value is None and len(game.possible_moves(cur_state)) == 0:
                value = 0.0

        return value, cur_state, cur_player, states, actions

    def is_leaf(self, state_int):
        # todo 为什么state_init不在self.probs中就是叶子节点
        # todo 书上说最终的游戏状态不会添加进统计中，在哪里控制的？
        # 难道叶子节点的定义就是，如果没有访问过就是叶子节点？
        return state_int not in self.probs

    def search_batch(self, count, batch_size, state_int, player, net, device="cpu"):
        '''
        mcts搜索 包含多次搜索，每次搜索包括：查找树的叶子、扩展叶子，进行备份
        todo 什么是扩展叶子
        param count: todo 搜索次数
        param batch_size: todo 批次大小
        param state_int: 当前状态
        param player: 当前玩家的索引
        param net: 当前玩家的神经网络
        param device: 训练设备

        '''
        # 为了能够更加有效的搜索扩展叶子，需要进行多次搜索，然后统一到单个NN中进行扩展
        # todo 解答P365的疑问
        # 执行多次搜索，每次搜索包括：查找树的叶子、扩展叶子，进行备份
        for _ in range(count):
            self.search_minibatch(batch_size, state_int, player, net, device)

    def search_minibatch(self, count, state_int, player, net, device="cpu"):
        """
        Perform several MCTS searches.
        执行MCTS搜搜

        param count: 进行棋盘下棋预测搜索的次数
        """
        # 存储搜索过程中，发现的最终状态的游戏过程或者扩展状态的游戏过程
        # [(当前状态下的价值，游戏状态过程，动作过程)]
        backup_queue = []
        # 探索到叶子状态，将叶子状态添加到expand_queue中，用于扩展
        expand_states = []
        # 探索到叶子状态，将叶子状态的玩家添加到expand_players中，用于扩展
        expand_players = []
        # 保存搜索过程中，发现的未被探索的状态，用于扩展（叶子状态、到达叶子状态的状态列表、每一个状态转换的执行动作）
        expand_queue = []
        # 保存搜索过程中，探索到的叶子节点状态（单个状态，不是状态序列）
        planned = set()
        for _ in range(count):
            # 进行一次搜索，直到游戏结束或者没有发现未被探索的状态
            value, leaf_state, leaf_player, states, actions = self.find_leaf(state_int, player)
            if value is not None:
                # 搜索发现的是最终状态，则保存至backup_queue
                backup_queue.append((value, states, actions))
            else:
                # 搜索发现的是未被探索的状态，并且该叶子状态是未被探索过的，则保存至expand_queue
                # 需要对叶子节点进行扩展
                if leaf_state not in planned:
                    # 将叶子节点添加到已经探索过的状态中
                    planned.add(leaf_state)
                    # 将int类型的状态转换为 bits list类型的状态
                    leaf_state_lists = game.decode_binary(leaf_state)
                    expand_states.append(leaf_state_lists)
                    expand_players.append(leaf_player)
                    expand_queue.append((leaf_state, states, actions))

        # do expansion of nodes
        if expand_queue:
            # 如果待扩展的状态不为空，则需要进行扩展
            # 将带扩展的状态转换为batch
            batch_v = model.state_lists_to_batch(expand_states, expand_players, device)
            # 将batch输入到神经网络中，拿到预测的动作以及奖励回报
            logits_v, values_v = net(batch_v)
            probs_v = F.softmax(logits_v, dim=1)
            # [:, 0]实在降维，因为原来的是二维的，所以降维后，就是一维的
            values = values_v.data.cpu().numpy()[:, 0]
            probs = probs_v.data.cpu().numpy()

            # create the nodes
            for (leaf_state, states, actions), value, prob in zip(expand_queue, values, probs):
                # [0] * game.GAME_COLS 创建一创建一个当前状态下，每列被访问的次数存储列表
                # 0表示每个动作被访问的0次
                self.visit_count[leaf_state] = [0] * game.GAME_COLS
                # [0.0] * game.GAME_COLS 创建一个当前状态下，每列的奖励回报存储列表，0表示每个动作的奖励回报为0
                self.value[leaf_state] = [0.0] * game.GAME_COLS
                # [0.0] * game.GAME_COLS 创建一个当前状态下，每列的奖励回报存储列表，0表示每个动作的奖励回报为0
                self.value_avg[leaf_state] = [0.0] * game.GAME_COLS
                # 将当前状态下，每列的动作概率存储到self.probs中
                self.probs[leaf_state] = prob
                # 将扩展后的状态添加到backup_queue中，用于后续的备份
                backup_queue.append((value, states, actions))

        # perform backup of the searches
        # todo 执行备份搜索
        for value, states, actions in backup_queue:
            # leaf state is not stored in states and actions, so the value of the leaf will be the value of the opponent
            # 由于最终状态以及叶子状态不存储在状态和动作中，所以对于最终状态下的值将是对手的值
            # 没有对应的状态和动作 故states[::-1], actions[::-1]要排除最后一个元素
            # 并且value要取反，因为对手赢，那么当前玩家就输了；对手输，那么当前玩家就赢了
            cur_value = -value
            # 从开始状态遍历到最终的状态
            for state_int, action in zip(states[::-1], actions[::-1]):
                # 统计当前状态下，执行的动作被访问的次数
                self.visit_count[state_int][action] += 1
                # 统计当前状态下，执行的动作的奖励回报
                self.value[state_int][action] += cur_value
                # 计算当前状态下，执行的动作的平均奖励回报
                self.value_avg[state_int][action] = self.value[state_int][action] / self.visit_count[state_int][action]
                # 再次取反，因为当前状态是对手的状态，而如果当前玩家是赢的状态
                # 那么对于对手来说，整个下棋的过程就是输的，回报都应该是对手回报取反
                cur_value = -cur_value

    def get_policy_value(self, state_int, tau=1):
        """
        todo 作用
        Extract policy and action-values by the state
        :param state_int: state of the board 提取当前状态下的动作概率以及奖励回报
        :param tau: tau=0 selects the most probable moves todo 作用
        :return: (probs, values)，每个动作的执行概率，每个动作的奖励回报
        """
        # 获取当前状态下，每列的动作被访问的次数
        counts = self.visit_count[state_int]
        if tau == 0:
            # 如果tau=0，则选择最可能的动作，则将访问次数最多的动概率设置为1，其他动作概率为0
            probs = [0.0] * game.GAME_COLS
            probs[np.argmax(counts)] = 1.0
        else:
            # 如果tau不为0，则选择概率最大的动作，概率越大，被选中的概率越大
            # 则遍历每一个动作，每个动作进行1/tau次方的计算
            counts = [count ** (1.0 / tau) for count in counts]
            # 计算和
            total = sum(counts)
            # 归一化为概率分布
            probs = [count / total for count in counts]
        values = self.value_avg[state_int]
        return probs, values
